/**
 * PHOSIDE: PHosphorylation Site IDentification Engine.
 * Copyright 2009 Digital Biology Lab, University of Missouri.
 * This library is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation; either version 3 of the License, or (at your option) any later
 * version. <p/> This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the License for more
 * details.
 */

package phoside.miscellaneous.test;

import java.io.IOException;

import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.TreeSet;

import phoside.PhosphoProteins;
import phoside.PhosphoProteinsImpl;

import phoside.Protein;
import phoside.ProteinFilter;

import phoside.data.util.PhosphoDataUtil;

import phoside.io.xml.PhosideXmlProteinsReader;
import phoside.io.xml.PhosideXmlProteinsWriter;
import phoside.io.xml.PhosphoProteinFieldValueFormatter;

import phoside.util.SamplingUtil;
import phoside.util.StaticFinalValues;

/**
 *
 * @author gjj
 */
public class TrainTestBuilder {

    public void build(String oriXml, String trainXml, String testXml, Set<PhosphoProteins.PhosphoType> types,
            int nTestPhos, int nTestNonPhos) throws Exception {
        // read
        PhosphoProteins proteins = new PhosphoProteinsImpl();
        PhosideXmlProteinsReader reader = new PhosideXmlProteinsReader(oriXml,
                proteins, new PhosphoProteinFieldValueFormatter());
        reader.read();

        // seperate into pos and neg
        PhosphoProteins proteins_pos = new PhosphoProteinsImpl();
        PhosphoProteins proteins_neg = new PhosphoProteinsImpl();

        Set<Character> aas = PhosphoProteins.PhosphoType.phosphorylatableAminoAcids(types);

        Iterator<Protein> it = proteins.proteinIterator();
        while (it.hasNext()) {
            Protein protein = it.next();
            String seq = protein.getSequence();
            String acc = protein.getAccession();
            Set<Integer> sites = proteins.getSites(acc);
            if (sites==null) {
                proteins_neg.addProtein(protein);
            } else {
                Set<Integer> newSites = new TreeSet();
                for (int site : sites) {
                    if (aas.contains(seq.charAt(site))) {
                        newSites.add(site);
                    }
                }
                protein.putInfo(StaticFinalValues.SITES, newSites);
                if (newSites.isEmpty()) {
                    proteins_neg.addProtein(protein);
                } else {
                    proteins_pos.addProtein(protein);
                }
            }
        }

        // sample
        PhosphoProteins proteins_train = new PhosphoProteinsImpl();
        PhosphoProteins proteins_test = new PhosphoProteinsImpl();

        int npos = proteins_pos.proteinCount();
        Set<Integer> idx_test = new HashSet(SamplingUtil.resampleWithoutReplacement(npos, nTestPhos));
        it = proteins_pos.proteinIterator();
        int count = 0;
        while (it.hasNext()) {
            Protein protein = it.next();
            if (idx_test.contains(count)) {
                proteins_test.addProtein(protein);
            } else {
                proteins_train.addProtein(protein);
            }
            count++;
        }

        int nneg = proteins_neg.proteinCount();
        idx_test = new HashSet(SamplingUtil.resampleWithoutReplacement(nneg, nTestNonPhos));
        it = proteins_neg.proteinIterator();
        count = 0;
        while (it.hasNext()) {
            Protein protein = it.next();
            if (idx_test.contains(count)) {
                proteins_test.addProtein(protein);
            } else {
                proteins_train.addProtein(protein);
            }
            count++;
        }


        // write
        PhosideXmlProteinsWriter writer = new PhosideXmlProteinsWriter(trainXml,
                new PhosphoProteinFieldValueFormatter());
        writer.write(proteins_train);

        writer = new PhosideXmlProteinsWriter(testXml,
                new PhosphoProteinFieldValueFormatter());
        writer.write(proteins_test);

    }

    public static void main(String[] args) {
        String app = "-ST";
        String outDir = "exp\\test1\\mouse"+app+"\\";
        String oriDir = "data\\phosphoelm\\";
        String oriXml = "Mouse-SwissProt-V57.8-20090922-PhosphoELM-V8.2-April2009-nr50";
//        String outDir = "exp\\test1\\Arabidopsis"+app+"\\";
//        String oriDir = "data\\phosphat\\";
//        String oriXml = "Arabidopsis-TAIR9-PhosPhAt3-nonambiguous-with-disorder-nr50";
//        String oriDir = "data\\phosphopep\\";
//        String oriXml = "Worm-UniProt-SwissProt-V15.8-20090922-PhophoPep-20091002-nr50";
        
        int nTestPhos = 100;
        int nTestNonPhos = 100;
        Set<PhosphoProteins.PhosphoType> types = new HashSet();
        types.add(PhosphoProteins.PhosphoType.SER);
        types.add(PhosphoProteins.PhosphoType.THR);
//        types.add(PhosphoProteins.PhosphoType.TYR);
        TrainTestBuilder builder = new TrainTestBuilder();
        try {
            builder.build(oriDir+oriXml+".xml", outDir+oriXml+"-train"+app+".xml", outDir+oriXml+"-test"+app+".xml", types, nTestPhos, nTestNonPhos);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
